#include <algorithm>
#include <iostream>
#include <iterator>
#include <stdexcept>
#include <type_traits>
#include <vector>

#include "tinyndarray.h"
#undef __EXCEPTIONS
#include "utest.h"
#include "cpptrace/cpptrace.hpp"
using namespace tinyndarray;

UTEST(ndarray, basic) {
    // 创建float类型数组
    auto float_arr = NdArray::Zeros(Dtype(FP32), {3, 3});
    float_arr[0] = 1.5f;       // 正确赋值
    float_arr[{1, 1}] = 2.5f;  // 多维索引

    try {
        float_arr[0] = 10.f;     // 自动转换为float (1.5f -> 10.0f)
        int val = float_arr[0];  // 抛出异常（类型不匹配）
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
    }

    // 创建int类型数组
    NdArray int_arr(Dtype(INT32), {4});
    for (size_t i = 0; i < int_arr.size(); ++i) {
        int_arr[i] = static_cast<int>(i * i);
    }

    // 点积运算
    auto vec1 = NdArray::Ones(Dtype(FP32), {5});
    auto vec2 = NdArray::Ones(Dtype(FP32), {5});
    auto dot_result = vec1.dot(vec2);
    float dot_value = dot_result;  // 标量转换

    // 迭代器使用
    for (auto& elem : int_arr) {
        elem = elem * 2;  // 通过代理修改元素
    }

    for (auto& elem : int_arr) {
        std::cout << static_cast<int>(elem) << std::endl;
    }

    std::cout << int_arr << std::endl;
    // 类型转换
    NdArray scalar_arr(Dtype(INT32), {1}, 42);
    int value = scalar_arr;  // 转换为int
}

UTEST(ndarray, concat) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {6, 6});
    auto m2 = NdArray::Uniform(0.f, 10.f, {6, 6});
    auto m = Concatenate({m1, m2}, 1);
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
    std::cout << m << std::endl;
}

UTEST(ndarray, transpose) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {2, 2, 2});
    auto m2 = Transpose(m1, {0, 2, 1});
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
}

UTEST(ndarray, split) {
    auto m = NdArray::Uniform(0.f, 10.f, {6, 6});
    auto mv = Split(m, 2, 1);
    std::cout << m << std::endl;
    std::cout << mv.size() << std::endl;
    std::cout << mv[0] << std::endl;
    std::cout << mv[1] << std::endl;
}

UTEST(ndarray, add) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {3, 3});
    auto m2 = NdArray::Uniform(0.f, 10.f, {1, 3});
    auto m = m1 + m2;
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
    std::cout << m << std::endl;
    auto m3 = m1 + 1.f;
    std::cout << m3 << std::endl;
}

UTEST(ndarray, cvt) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {3, 3});
    auto m2 = m1.astype(Dtype(INT32));
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
}

UTEST(ndarray, reduce) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {3, 3});
    auto m2 = m1.sum({1});
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
}

UTEST(ndarray, neg) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {3, 3});
    auto m2 = -m1;
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
}

UTEST(ndarray, io) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {3, 3});
    m1.Save("m1.npy");
    auto m2 = Load("m1.npy");
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
}

UTEST(ndarray, init) {
    NdArray float_arr = {{{-0.4f, 0.3f}, {-0.2f, 0.1f}}, {{-0.1f, 0.2f}, {-0.3f, 0.4f}}};
    std::cout << float_arr << std::endl;
    NdArray int_arr = {{{-1, 2}, {-2, 3}}, {{1, 0}, {-3, 4}}};
    std::cout << int_arr << std::endl;
}

UTEST(ndarray, slice) {
    auto m = NdArray::Uniform(0.f, 10.f, {6, 6});
    std::cout << m << std::endl;
    auto m1 = m({0, 4}, {1, 5});
    auto one = NdArray::Ones(Dtype(FP32), {3, 3});
    m1({1, 3}, {2, 4}) = one({0, 2}, {1, 3}) + 1.f;
    m1[{2, 3}] = 3.f;
    std::cout << m << std::endl;
}

UTEST(ndarray, max) {
    auto m = NdArray::Uniform(0.f, 10.f, {6, 6});
    auto m1 = m.max({1});
    std::cout << m << std::endl;
    std::cout << m1 << std::endl;
}

UTEST(ndarray, broadcast) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {3, 3});
    auto m2 = NdArray::Uniform(0.f, 10.f, {3, 1});
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
    auto m3 = m1 + m2;
    std::cout << m3 << std::endl;
}

UTEST(ndarray, cmp) {
    auto m1 = NdArray::Uniform(0.f, 10.f, {3, 3});
    auto m2 = NdArray::Uniform(-1.f, 11.f, {3, 3});
    std::cout << m1 << std::endl;
    std::cout << m2 << std::endl;
    std::cout << (m1 < m2) << std::endl;
}

UTEST(ndarray, softmax) {
    auto x = NdArray::Uniform(0.f, 10.f, {3, 3});
    auto e_x = Exp(x - x.max({-1}, true));
    auto s = e_x / e_x.sum({-1}, true);
    std::cout << x << std::endl;
    std::cout << s << std::endl;
}

UTEST(ndarray, fp16) {
    auto x = NdArray::Uniform(float16_t(1.f), float16_t(10.f), {3, 3});
    std::cout << x << std::endl;
    auto xf = x.astype(Dtype(FP32));
    xf = xf * 2.f;
    x = xf.astype(Dtype(FP16));
    std::cout << x << std::endl;
}

UTEST(ndarray, quant) {
    auto x = NdArray::Uniform(-10.f, 10.f, {3, 3});
    std::cout << "===================量化前的矩阵===================\n";
    std::cout << x << std::endl;
    auto qx = x.astype(Dtype(QINT8, {2.0, 10})); // 2.0是scale，10是zeropoint
    std::cout << "===================量化后的矩阵===================\n";
    std::cout << qx << std::endl;
    x = qx.astype(Dtype(FP32));
    std::cout << "==================反量化回来的矩阵=================\n";
    std::cout << x << std::endl;
}

UTEST_MAIN()